//
// Created by SongpingWang on 2025/2/8.
//

#include "dataset.h"


CustomDataset::CustomDataset(const std::vector<std::string>& filenames,
                             int input_size,
                             const std::map<std::string, double>& params,
                             bool augment)
        : filenames_(filenames),
          input_size_(input_size),
          params_(params),
          augment_(augment) {

    // Read labels
    labels_ = load_label(filenames_);
    n_ = labels_.size();

    indices_.resize(n_);
    std::iota(indices_.begin(), indices_.end(), 0);

    // Albumentations (optional, only used if package is installed)
    albumentations_ = Albumentations();
}


torch::data::Example<> CustomDataset::get(size_t index) {
    std::lock_guard<std::mutex> lock(mtx_); // Ensure thread safety

    index = indices_[index];

    cv::Mat image;
    torch::Tensor label;

    if (params_.at("mosaic") > 0 && static_cast<double>(rand()) / RAND_MAX < params_.at("mosaic")) {
        // Load MOSAIC
        std::tie(image, label) = load_mosaic(index);

        // MixUp augmentation
        if (static_cast<double>(rand()) / RAND_MAX < params_.at("mix_up")) {
            size_t mix_index = indices_[rand() % n_];
            auto [mix_image1, mix_label1] = std::make_tuple(image, label);
            auto [mix_image2, mix_label2] = load_mosaic(mix_index);

            std::tie(image, label) = mix_up(mix_image1, mix_label1, mix_image2, mix_label2);
        }
    } else {
        // Load image
        std::tie(image, std::ignore) = load_image(index);
        auto [h, w] = std::make_pair(image.rows, image.cols);

        // Resize
        std::tie(image, std::ignore, std::ignore) = resize(image, input_size_, augment_);

        label = labels_[index].clone();
        if (label.size(0) > 0) {
            label.slice(1, 1, 5) = wh2xy(label.slice(1, 1, 5), w, h, 0, 0);
        }
        if (augment_) {
            std::tie(image, label) = random_perspective(image, label, params_);
        }
    }

    auto nl = label.size(0);  // number of labels
    auto [h, w] = std::make_pair(image.rows, image.cols);
    auto cls = label.slice(1, 0, 1);
    auto box = xy2wh(label.slice(1, 1, 5), w, h);

    if (augment_) {
        // Convert torch::Tensor to std::vector for Albumentations
        std::vector<std::vector<float>> box_vec(box.size(0), std::vector<float>(4));
        std::vector<int> cls_vec(cls.size(0));
        for (int i = 0; i < box.size(0); ++i) {
            box_vec[i] = {box[i][0].item<float>(), box[i][1].item<float>(), box[i][2].item<float>(), box[i][3].item<float>()};
            cls_vec[i] = static_cast<int>(cls[i].item<float>());
        }

        // Albumentations
        std::tie(image, box_vec, cls_vec) = albumentations_(image, box_vec, cls_vec);
        nl = box_vec.size();  // update after albumentations

        // Convert back to torch::Tensor
        box = torch::zeros({nl, 4});
        cls = torch::zeros({nl, 1}, torch::kInt64);
        for (int i = 0; i < nl; ++i) {
            box[i] = torch::tensor(box_vec[i], torch::kFloat);
            cls[i] = cls_vec[i];
        }

        // HSV color-space
        augment_hsv(image, params_);
        // Flip up-down
        if (static_cast<double>(rand()) / RAND_MAX < params_.at("flip_ud")) {
            cv::flip(image, image, 0);
            if (nl) {
                box.select(1, 1) = 1 - box.select(1, 1);
            }
        }
        // Flip left-right
        if (static_cast<double>(rand()) / RAND_MAX < params_.at("flip_lr")) {
            cv::flip(image, image, 1);
            if (nl) {
                box.select(1, 0) = 1 - box.select(1, 0);
            }
        }
    }

    torch::Tensor target_cls = torch::zeros({nl, 1});
    torch::Tensor target_box = torch::zeros({nl, 4});
    if (nl) {
        target_cls = cls;
        target_box = box;
    }

    // Convert HWC to CHW, BGR to RGB
    cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
    auto sample = torch::from_blob(image.data, {1, image.rows, image.cols, 3}, torch::kByte).to(torch::kFloat).permute({0, 3, 1, 2}).contiguous();

    return {sample, target_cls, target_box, torch::zeros(nl)};
}



std::tuple<cv::Mat, std::pair<int, int>> CustomDataset::load_image(size_t index) {
    cv::Mat image = cv::imread(filenames_[index]);
    int h = image.rows;
    int w = image.cols;
    double r = static_cast<double>(input_size_) / std::max(h, w);
    if (r != 1) {
        cv::resize(image, image, cv::Size(static_cast<int>(w * r), static_cast<int>(h * r)), 0, 0, augment_ ? resample() : cv::INTER_LINEAR);
    }
    return std::make_tuple(image, std::make_pair(h, w));
}


std::tuple<cv::Mat, torch::Tensor> CustomDataset::load_mosaic(size_t index) {
    std::vector<torch::Tensor> label4;
    int border = -input_size_ / 2;
    int double_input_size = 2 * input_size_;
    cv::Mat image4 = cv::Mat::zeros(double_input_size, double_input_size, CV_8UC3);

    // Seed with a real random value, if available
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0, double_input_size + border * 2 - 1);
    std::uniform_int_distribution<> index_dis(0, (int)n_ - 1);

    int xc = dis(gen);
    int yc = dis(gen);

    std::vector<size_t> indices = {index};
    for (int i = 0; i < 3; ++i) {
        indices.push_back(index_dis(gen));
    }
    std::shuffle(indices.begin(), indices.end(), gen);

    for (size_t i = 0; i < 4; ++i) {
        size_t idx = indices[i];
        cv::Mat image;
        std::tie(image, std::ignore) = load_image(idx);
        int h = image.rows;
        int w = image.cols;

        int x1a, y1a, x2a, y2a, x1b, y1b, x2b, y2b;
        if (i == 0) {  // top left
            x1a = std::max(xc - w, 0);
            y1a = std::max(yc - h, 0);
            x2a = xc;
            y2a = yc;
        } else if (i == 1) {  // top right
            x1a = xc;
            y1a = std::max(yc - h, 0);
            x2a = std::min(xc + w, double_input_size);
            y2a = yc;
        } else if (i == 2) {  // bottom left
            x1a = std::max(xc - w, 0);
            y1a = yc;
            x2a = xc;
            y2a = std::min(double_input_size, yc + h);
        } else {  // bottom right
            x1a = xc;
            y1a = yc;
            x2a = std::min(xc + w, double_input_size);
            y2a = std::min(double_input_size, yc + h);
        }

        x1b = (i % 2 == 0) ? w - (x2a - x1a) : 0;
        y1b = (i < 2) ? h - (y2a - y1a) : 0;
        x2b = (i % 2 == 0) ? w : std::min(w, x2a - x1a);
        y2b = (i < 2) ? h : std::min(h, y2a - y1a);

        int pad_w = x1a - x1b;
        int pad_h = y1a - y1b;
        image(cv::Rect(x1b, y1b, x2b - x1b, y2b - y1b)).copyTo(image4(cv::Rect(x1a, y1a, x2a - x1a, y2a - y1a)));

        // Labels
        auto label = labels_[idx].clone();
        if (label.size(0) > 0) {
            label.slice(1, 1, 5) = wh2xy(label.slice(1, 1, 5), w, h, pad_w, pad_h);
        }
        label4.push_back(label);
    }

    // Concat/clip labels
    auto label4_tensor = torch::cat(label4, 0);
    label4_tensor.slice(1, 1, 5) = torch::clamp(label4_tensor.slice(1, 1, 5), 0, double_input_size);

    // Augment
    std::tie(image4, label4_tensor) = random_perspective(image4, label4_tensor, params_, std::make_pair(border, border));

    return std::make_tuple(image4, label4_tensor);
}


std::unordered_map<std::string, torch::Tensor> CustomDataset::load_label(const std::vector<std::string>& filenames) {
    std::unordered_map<std::string, torch::Tensor> labels;
    for (const auto& filename : filenames) {
        try {
            cv::Mat image = cv::imread(filename);
            if (image.empty()) continue;

            int h = image.rows;
            int w = image.cols;
            assert(h > 9 && w > 9);

            std::string label_path = filename;
            size_t pos = label_path.find("images");
            if (pos != std::string::npos) {
                label_path.replace(pos, 6, "labels");
                label_path.replace(label_path.find(".jpg"), 4, ".txt");
            }

            std::ifstream file(label_path);
            std::vector<std::vector<float>> label_data;
            std::string line;
            while (std::getline(file, line)) {
                std::istringstream ss(line);
                std::vector<float> values((std::istream_iterator<float>(ss)), std::istream_iterator<float>());
                label_data.push_back(values);
            }

            torch::Tensor label = torch::zeros({static_cast<long>(label_data.size()), 5});
            for (int i = 0; i < label_data.size(); ++i) {
                label[i] = torch::tensor(label_data[i]);
            }

            labels[filename] = label;
        } catch (...) {
            continue;
        }
    }
    return labels;
}
